import os
import torch
import logging
from tqdm import tqdm
from dataloader_my import build_dataloader
from model_my_fpn import ResNetConvLSTM
from log_manage import logger_init


def test_model():
    test_dataset_dir = r'/home/xian/mzs_project/Convlstm/dataset_fog/test'
    experiment_name = "name1"
    log_dir = os.path.join("logs", experiment_name)
    os.makedirs(log_dir, exist_ok=True)
    save_dir = os.path.join("./save_model", experiment_name)
    model_path = os.path.join(save_dir, "best_model.pt")

    output_log_file = os.path.join("./test_result", experiment_name, "test_results.txt")
    output_log_dir = os.path.dirname(output_log_file)
    os.makedirs(output_log_dir, exist_ok=True)
    frame_len = 15
    image_size = (360, 640)
    batch_size = 1

    num_classes = 2
    pretrained_path = None
    categories = ["nofoggy", "foggy"]

    logger_init(log_file_name='test_log', log_level=logging.INFO, log_dir=log_dir)
    device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
    logging.info(f"Using device: {device}")

    logging.info(f"加载测试数据从: {test_dataset_dir}")
    test_loader = build_dataloader(
        root_dir=test_dataset_dir,
        frame_len=frame_len,
        batch_size=batch_size,
        num_workers=0,
        image_size=image_size,
        mode='test'
    )

    logging.info(f"加载模型: {model_path}")
    model = ResNetConvLSTM(num_classes=num_classes, pretrain_path=pretrained_path).to(device)

    if not os.path.exists(model_path):
        logging.error(f"模型权重文件不存在: {model_path}")
        return

    model.load_state_dict(torch.load(model_path, map_location=device))
    model.eval()

    dataset = test_loader.dataset
    video_infos = dataset.video_infos

    correct, total = 0, 0

    logging.info(f"开始测试... 结果将保存到 {output_log_file}")
    with open(output_log_file, 'w', encoding='utf-8') as f:
        f.write("测试结果日志 (状态 | 视频目录 | 起始帧 | 真实标签 | 预测标签)\n")
        f.write("-" * 80 + "\n")

        progress = tqdm(test_loader, desc="Testing")
        for i, (x_batch, y_batch) in enumerate(progress):

            video_dir_key, start_idx, true_label_int = video_infos[i]

            x_batch, y_batch = x_batch.to(device), y_batch.to(device)

            with torch.no_grad():
                logits = model(x_batch)
            pred_label_int = logits.argmax(1).item()
            true_label_from_loader = y_batch.item()

            assert true_label_int == true_label_from_loader, "Dataloader 标签和 video_infos 不匹配!"

            total += 1
            if pred_label_int == true_label_int:
                correct += 1

            true_label_str = categories[true_label_int]
            pred_label_str = categories[pred_label_int]

            status_tag = "[CORRECT]" if pred_label_int == true_label_int else "[WRONG]"

            log_line = f"{status_tag:<9} | {video_dir_key} | {start_idx:<8} | {true_label_str:<7} | {pred_label_str}\n"
            f.write(log_line)

    accuracy = (correct / total) * 100 if total > 0 else 0
    logging.info("--- 测试完成 ---")
    logging.info(f"总样本数: {total}")
    logging.info(f"正确数: {correct}")
    logging.info(f"准确率: {accuracy:.2f}%")
    logging.info(f"详细日志已保存到: {output_log_file}")
    print(f"\n测试完成! 准确率: {accuracy:.2f}%")
    print(f"详细日志已保存到: {output_log_file}")


if __name__ == "__main__":
    test_model()
